package com.shure.modules.llm.controller;

import com.shure.modules.llm.model.MessageDto;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

@RestController
@RequestMapping("/ai")
public class ChatController {

    private final ChatModel chatModel;
    private final ToolCallbackProvider toolCallbackProvider;

    public ChatController(OllamaChatModel chatModel, ToolCallbackProvider toolCallbackProvider) {
        this.chatModel = chatModel;
        this.toolCallbackProvider = toolCallbackProvider;
    }

    @PostMapping("/chat")
    public ResponseEntity<String> chat(@RequestBody MessageDto dto) {
        Prompt prompt = this.buildPrompt(dto);
        return ResponseEntity.ok(this.chatModel.call(prompt).getResult().getOutput().getText());
    }

    @PostMapping(value = "/chat/stream", produces = "text/event-stream")
    public Flux<ChatResponse> chatStream(@RequestBody MessageDto dto) {
        return this.chatModel.stream(this.buildPrompt(dto));
    }

    private Prompt buildPrompt(MessageDto dto) {
        List<Message> chatMessages = new ArrayList<>();
        for (MessageDto.Message message : dto.getMessages()) {
            if (message.getRole().equals("user")) {
                chatMessages.add(new UserMessage(message.getContent()));
            }
            if (message.getRole().equals("assistant")) {
                chatMessages.add(new AssistantMessage(message.getContent()));
            }
            if (message.getRole().equals("system")) {
                chatMessages.add(new SystemMessage(message.getContent()));
            }
        }
        return new Prompt(chatMessages);
    }

    @GetMapping("/mcp")
    public String generate(@RequestParam(value = "message") String message) {
        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultTools(toolCallbackProvider.getToolCallbacks())
                .build();
        return chatClient.prompt(message).call().content();
    }

    @GetMapping("/mcp/stream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message") String message) {
        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultTools(toolCallbackProvider.getToolCallbacks())
                .build();
        return chatClient.prompt(message).stream().chatResponse();
    }


}
